import colorsys
import warnings
from typing import Sequence

import cv2
import numpy as np
import torch
from PIL import Image
from scipy.ndimage import gaussian_filter, map_coordinates
from skimage.transform import rotate
from torch import Tensor
from torchvision import transforms
from torchvision.transforms import functional as F, InterpolationMode
from torchvision.transforms.functional import _interpolation_modes_from_int


def normalize(x, epsilon=1e-7, axis=(1, 2)):
    x -= np.mean(x, axis=axis, keepdims=True)
    x /= np.std(x, axis=axis, keepdims=True) + epsilon
    return x


def letterbox_image_mask(image, mask, size):
    mask = Image.fromarray(np.array(mask))
    w, h = size
    '''resize image with unchanged aspect ratio using padding'''
    new_image, nw, nh = letterbox_image(image, size)

    mask = mask.resize((nw, nh), Image.NEAREST)
    new_label = Image.new('L', size, 0)
    new_label.paste(mask, ((w - nw) // 2, (h - nh) // 2))

    return new_image / 255, new_label


def letterbox_image(image, size, image_type='RGB'):
    iw, ih = image.size
    w, h = size
    scale = min(w / iw, h / ih)
    nw = int(iw * scale)
    nh = int(ih * scale)

    image = image.resize((nw, nh), Image.BICUBIC)
    new_image = Image.new(image_type, size)
    new_image.paste(image, ((w - nw) // 2, (h - nh) // 2))
    return new_image, nw, nh


def rand(a=0, b=1):
    return np.random.rand() * (b - a) + a


def deeplab_dataset_collate(batch):
    images = []
    pngs = []
    seg_labels = []
    for img, png, labels in batch:
        images.append(img)
        pngs.append(png)
        seg_labels.append(labels)
    images = np.array(images)
    pngs = np.array(pngs)
    seg_labels = np.array(seg_labels)
    return images, pngs, seg_labels


def random_elastic_deformation(image,
                               alpha=500,
                               sigma=20,
                               mode='nearest',
                               random_state=None):
    """Elastic deformation of images as described in [Simard2003]_.
    .. [Simard2003] Simard, Steinkraus and Platt, "Best Practices for
       Convolutional Neural Networks applied to Visual Document Analysis", in
       Proc. of the International Conference on Document Analysis and
       Recognition, 2003.
    """
    assert len(image.shape) == 3

    if random_state is None:
        random_state = np.random.RandomState(None)

    height, width, channels = image.shape

    dx = gaussian_filter(2 * random_state.rand(height, width) - 1,
                         sigma, mode="constant", cval=0) * alpha
    dy = gaussian_filter(2 * random_state.rand(height, width) - 1,
                         sigma, mode="constant", cval=0) * alpha

    x, y = np.meshgrid(np.arange(height), np.arange(width), indexing='ij')
    indices = (np.repeat(np.ravel(x + dx), channels),
               np.repeat(np.ravel(y + dy), channels),
               np.tile(np.arange(channels), height * width))

    values = map_coordinates(image, indices, order=1, mode=mode)

    return values.reshape((height, width, channels))


def maybe_rotate(image):
    # orient image in landscape
    height, width = image.shape
    return np.rot90(image) if width < height else image


# one-hot encode
def to_categorical(y, num_classes=None, dtype='long'):
    y = np.array(y, dtype='int')
    input_shape = y.shape
    if input_shape and input_shape[-1] == 1 and len(input_shape) > 1:
        input_shape = tuple(input_shape[:-1])
    y = y.ravel()
    if not num_classes:
        num_classes = np.max(y) + 1
    n = y.shape[0]
    categorical = np.zeros((n, num_classes), dtype=dtype)
    categorical[np.arange(n), y] = 1
    output_shape = input_shape + (num_classes,)
    categorical = np.reshape(categorical, output_shape)
    return categorical


def get_random_data(image, mask, input_shape, jitter=.3, hue=.1, sat=1.5, val=1.5):
    mask = Image.fromarray(np.array(mask))

    h, w = input_shape
    # resize image
    rand_jit1 = rand(1 - jitter, 1 + jitter)
    rand_jit2 = rand(1 - jitter, 1 + jitter)
    new_ar = w / h * rand_jit1 / rand_jit2

    scale = rand(0.5, 1.5)
    if new_ar < 1:
        nh = int(scale * h)
        nw = int(nh * new_ar)
    else:
        nw = int(scale * w)
        nh = int(nw / new_ar)
    image = image.resize((nw, nh), Image.BICUBIC)
    mask = mask.resize((nw, nh), Image.NEAREST)
    mask = mask.convert("L")

    # flip image or not
    flip = rand() < .5
    if flip:
        image = image.transpose(Image.FLIP_LEFT_RIGHT)
        mask = mask.transpose(Image.FLIP_LEFT_RIGHT)

    # place image
    dx = int(rand(0, w - nw))
    dy = int(rand(0, h - nh))
    new_image = Image.new('RGB', (w, h), (128, 128, 128))
    new_label = Image.new('L', (w, h), (0))
    new_image.paste(image, (dx, dy))
    new_label.paste(mask, (dx, dy))
    image = new_image
    mask = new_label

    # distort image
    hue = rand(-hue, hue)
    sat = rand(1, sat) if rand() < .5 else 1 / rand(1, sat)
    val = rand(1, val) if rand() < .5 else 1 / rand(1, val)
    x = cv2.cvtColor(np.array(image, np.float32) / 255, cv2.COLOR_RGB2HSV)
    x[..., 0] += hue * 360
    x[..., 0][x[..., 0] > 1] -= 1
    x[..., 0][x[..., 0] < 0] += 1
    x[..., 1] *= sat
    x[..., 2] *= val
    x[x[:, :, 0] > 360, 0] = 360
    x[:, :, 1:][x[:, :, 1:] > 1] = 1
    x[x < 0] = 0
    image_data = cv2.cvtColor(x, cv2.COLOR_HSV2RGB)
    return image_data, mask


def gray2colorfulRgb(num_classes, label, orininal_w, orininal_h):
    if num_classes <= 21:
        colors = [(0, 0, 0), (128, 0, 0), (0, 128, 0), (128, 128, 0), (0, 0, 128), (128, 0, 128),
                  (0, 128, 128),
                  (128, 128, 128), (64, 0, 0), (192, 0, 0), (64, 128, 0), (192, 128, 0), (64, 0, 128),
                  (192, 0, 128),
                  (64, 128, 128), (192, 128, 128), (0, 64, 0), (128, 64, 0), (0, 192, 0), (128, 192, 0),
                  (0, 64, 128), (128, 64, 12)]
    else:
        # set different colors for the picture frame
        # hsv_tuples = [(x / len(self.class_names), 1., 1.)
        #               for x in range(len(self.class_names))]
        hsv_tuples = [(x / num_classes, 1., 1.)
                      for x in range(num_classes)]
        colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
        colors = list(
            map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
                colors))

    seg_img = np.zeros((np.shape(label)[0], np.shape(label)[1], 3))
    for c in range(num_classes):
        seg_img[:, :, 0] += ((label[:, :] == c) * (colors[c][0])).astype('uint8')
        seg_img[:, :, 1] += ((label[:, :] == c) * (colors[c][1])).astype('uint8')
        seg_img[:, :, 2] += ((label[:, :] == c) * (colors[c][2])).astype('uint8')

    image = Image.fromarray(np.uint8(seg_img)).resize((orininal_w, orininal_h)).convert('RGB')

    return image


def gray2rgb(image):
    w, h = image.shape
    image += np.abs(np.min(image))
    image_max = np.abs(np.max(image))
    if image_max > 0:
        image /= image_max
    ret = np.empty((w, h, 3), dtype=np.uint8)
    ret[:, :, 2] = ret[:, :, 1] = ret[:, :, 0] = image * 255
    return ret


def arrayPaste(array, w, iw, h, ih):
    ins = np.zeros(((w - iw) // 2, array.shape[1]))
    array = np.insert(array, 0, ins, axis=0)
    ins = np.zeros((w - array.shape[0], array.shape[1]))
    array = np.insert(ins, 0, array, axis=0)
    ins = np.zeros((array.shape[0], (h - ih) // 2))
    array = np.insert(array, [0], ins, axis=1)
    ins = np.zeros((array.shape[0], h - array.shape[1]))
    array = np.insert(ins, [0], array, axis=1)
    return array


class Compose(transforms.Compose):
    def __call__(self, img, mask):
        for t in self.transforms:
            img, mask = t(img, mask)
        return img, mask


class ColorJitter(transforms.ColorJitter):

    def forward(self, img, mask):
        fn_idx, brightness_factor, contrast_factor, saturation_factor, hue_factor = \
            self.get_params(self.brightness, self.contrast, self.saturation, self.hue)

        for fn_id in fn_idx:
            if fn_id == 0 and brightness_factor is not None:
                img = F.adjust_brightness(img, brightness_factor)
            elif fn_id == 1 and contrast_factor is not None:
                img = F.adjust_contrast(img, contrast_factor)
            elif fn_id == 2 and saturation_factor is not None:
                img = F.adjust_saturation(img, saturation_factor)
            elif fn_id == 3 and hue_factor is not None:
                img = F.adjust_hue(img, hue_factor)
        return img, mask


class RandomHorizontalFlip(transforms.RandomHorizontalFlip):
    def forward(self, img, mask):
        if torch.rand(1) < self.p:
            return F.hflip(img), F.hflip(mask)
        return img, mask


class RandomVerticalFlip(transforms.RandomVerticalFlip):

    def forward(self, img, mask):
        if torch.rand(1) < self.p:
            return F.vflip(img), F.vflip(mask)
        return img, mask


class Resize(transforms.Resize):
    def forward(self, img, mask):
        return F.resize(img, self.size, self.interpolation, self.max_size, self.antialias), F.resize(mask, self.size,
                                                                                                     self.interpolation,
                                                                                                     self.max_size,
                                                                                                     self.antialias)


class Scale(torch.nn.Module):
    def __init__(self, size, interpolation=InterpolationMode.BILINEAR, max_size=None, antialias=None):
        super().__init__()
        if not isinstance(size, (int, Sequence)):
            raise TypeError("Size should be int or sequence. Got {}".format(type(size)))
        if isinstance(size, Sequence) and len(size) not in (1, 2):
            raise ValueError("If size is a sequence, it should have 1 or 2 values")
        self.size = size
        self.max_size = max_size

        # Backward compatibility with integer value
        if isinstance(interpolation, int):
            warnings.warn(
                "Argument interpolation should be of type InterpolationMode instead of int. "
                "Please, use InterpolationMode enum."
            )
            interpolation = _interpolation_modes_from_int(interpolation)

        self.interpolation = interpolation
        self.antialias = antialias

    def forward(self, img, mask):
        w, h = self.size
        iw, ih = img.size

        img_arr = np.asarray(img)
        mask_arr = np.asarray(mask)
        img_arr = arrayPaste(img_arr, h, ih, w, iw)
        mask_arr = arrayPaste(mask_arr, h, ih, w, iw)
        return Image.fromarray(img_arr).convert(img.mode), Image.fromarray(mask_arr).convert(mask.mode)


class RandomRotation(transforms.RandomRotation):
    def forward(self, img, mask):
        fill = self.fill
        if isinstance(img, Tensor):
            if isinstance(fill, (int, float)):
                fill = [float(fill)] * F._get_image_num_channels(img)
            else:
                fill = [float(f) for f in fill]
        angle = self.get_params(self.degrees)

        return F.rotate(img, angle, self.resample, self.expand, self.center, fill), F.rotate(mask, angle, self.resample,
                                                                                             self.expand, self.center,
                                                                                             fill)


class ToTensor(transforms.ToTensor):
    def __call__(self, img, mask):
        return F.to_tensor(img), mask


class Normalize(transforms.Normalize):

    def forward(self, img: Tensor, mask: np.ndarray) -> Tensor:
        return F.normalize(img, self.mean, self.std, self.inplace), mask
